
open Histo

let _ =
      Gsl_error.init ()

let r=Gsl_rng.make (Gsl_rng.default ())

;;
open Gsl_fun
let construct_fun ~(fit_f: float array -> float -> float) ~deriv ~data_x_y_yerr=
    let multi_f ~x ~f=
        let arr_par=Gsl_vector.to_array x in
(*        Printf.printf "pars=%f %f %f; f(4.)=%f\n" arr_par.(0) arr_par.(1)
                arr_par.(2) (fit_f arr_par 4.);*)
        let cur_f=fit_f arr_par in
        for i=0 to Gsl_vector.length f - 1 do
            let (xn,y,y_err) = data_x_y_yerr.(i) in
            f.{i} <- ((cur_f xn) -. y) /. y_err;
(*            Printf.printf "f(%i)=%f (cur_f(%f)=%f y=%f e=%f)\n" i f.{i} xn
             (cur_f xn) y y_err*)
        done
    in
    (* calculate Jacobian of fit_f over parameters, e.g. j is matrix
     * j(k,l) = df_k / dx_l *)
    let multi_df ~x ~j=
        let arr_par=Gsl_vector.to_array x in
        let unfixed f_k l xn x_l=(
            let a=Array.copy arr_par in
            a.(l) <- x_l;
            f_k a xn
            )
        in
        let partial_f_k_make (xn,y,y_err) l = 
            match deriv with
              Some  f -> 
                  (f arr_par xn ~by:l) /. y_err
            | None ->
                  let f_k arr_par xn = 
                       ((fit_f arr_par xn) -. y) /. y_err in
                  let {res=res;err=_} =
                      Gsl_diff.central (unfixed f_k l xn) arr_par.(l) in
                  res
        in
        let (n,p) = Gsl_matrix.dims j in
        for k=0 to n-1 do
            let partial_f_k l = partial_f_k_make data_x_y_yerr.(k) l in
            for l=0 to p-1 do
                j.{k, l} <- partial_f_k l
(*                let f_k arr_par xn= 
                    ((fit_f arr_par xn) -. y) /. y_err in
                j.{k, l} <- partial' f_k l xn;
                *)
(*                Printf.printf "J(%d,%d)=%f\n" k l j.{k, l}*)
            done
        done
    in
    let multi_fdf ~x ~f ~j =
        multi_f ~x ~f;
        multi_df ~x ~j
    in
    {   multi_f = multi_f;
        multi_df = multi_df;
        multi_fdf = multi_fdf
    }
;;

let print_state n p iter s =
    let x = Gsl_vector.create p in
    let f = Gsl_vector.create n in
    Gsl_multifit_nlin.get_state s ~x ~f () ;
    Printf.printf "iter: %3u " iter;
    Array.iteri (fun i par -> 
        Printf.printf " p(%i) = %15.8f ; " i par) (Gsl_vector.to_array x);
    Printf.printf "  |f(x)| = %g\n" (Gsl_blas.nrm2 f)
(*    Printf.printf "iter: %3u x = %15.8f % 15.8f % 15.8f |f(x)| = %g\n"
    iter x.{0} x.{1} x.{2} (Gsl_blas.nrm2 f)*)
;;

let maxiter = 500
let epsabs = 1e-4
let epsrel = 1e-4

let solve n x_init ~f=
(*    let n = Array.length y in*)
    let p = Array.length x_init in
    let s=
        Gsl_multifit_nlin.make
            Gsl_multifit_nlin.LMSDER ~n ~p f
            (Gsl_vector.of_array x_init)
    in
    let rec proc iter =
        Gsl_multifit_nlin.iterate s ;
        let status = Gsl_multifit_nlin.test_delta s ~epsabs ~epsrel in
        (match status with
        | true ->
(*                Printf.printf "\nstatus = converged, %i iterations\n" iter; *)
                ()
        | false when iter >= maxiter ->
                print_state n p iter s;
                Printf.printf "\nstatus = too many iterations\n";
                flush stdout;
        | false ->
                proc (succ iter)
        )
    in
    proc 1;
    let pos = Gsl_vector.create p in
    Gsl_multifit_nlin.position s pos ;
    let covar = Gsl_matrix.create p p in
    Gsl_multifit_nlin.covar s 0.0 covar;
    Gsl_vector.to_array pos, Gsl_matrix.to_arrays covar
;;

let construct_chi2 ~(f: float array -> float -> float) 
        ~deriv ~data_x_y_yerr ~start:arr_par n_par =
(*    let unfixed f_k l xn x_l=(
        let a=Array.copy arr_par in
      a.(l) <- x_l;
        f_k a xn
        )
    in*)
    let subst_par l p_l =
      let a = Array.copy arr_par in
	a.(l) <- p_l;
	a
    in
    let chi2 p =
        let f_p xn = f (subst_par n_par p) xn in
        Array.fold_left (fun sum (xn,y,y_err) ->
            sum +. ((f_p xn) -. y)**2. /. y_err**2.
        ) 0. data_x_y_yerr
(*        let sum = ref 0. in
        for i=0 to Array.length data_x_y_yerr - 1 do
            let (xn,y,y_err) = data_x_y_yerr.(i) in
            sum := !sum +. ((cur_f xn) -. y)**2. /. y_err**2.;
        done*)
    in
    let chi2_df p =
        match deriv with
            Some deriv_f ->
              let f_p xn = f (subst_par n_par p) xn in
              let deriv_f_p xn = deriv_f (subst_par n_par p) xn in
                Array.fold_left (fun sum (xn,y,y_err) ->
                    sum +. (deriv_f_p xn ~by:n_par) *. 
                            2. *.((f_p xn) -. y) /. y_err**2.
                ) 0. data_x_y_yerr
          | None ->
                let { res = res; err = _ } =
                    Gsl_diff.central chi2 p in
                res
    in
    let ndf = Array.length data_x_y_yerr in
    let chi2_minimum = chi2 arr_par.(n_par) in
    let chi2_f x = chi2 x -. chi2_minimum -. (float ndf) in
    let chi2_fdf x =
        chi2_f x, chi2_df x
    in
    {   f = chi2_f;
        df = chi2_df;
        fdf = chi2_fdf
    }
;;

let calc_chi2 x_y_e f=
    let chi2sum=ref 0.
    and ndf=ref 0 in
    Array.iteri (fun i (x,y,e) ->
        if y>0. 
        then (
            chi2sum := !chi2sum +. ((f x)-.y)**2. /. e**2.;
            incr ndf
        );
    ) x_y_e;
    (!chi2sum,!ndf)

let print_chi2 x_y_e f=
    let (chi2,ndf) = calc_chi2 x_y_e f in
    Printf.printf "chi2/ndf=%f/%i=%f, prob=%f\n" chi2 ndf 
        (chi2 /. float ndf) (Gsl_helpers.prob chi2 ndf)

let fit_lsq ~data_x_y_yerr ~by_function:f ~start_par ~deriv =
    let gsl_multi=construct_fun ~fit_f:f ~deriv ~data_x_y_yerr in
    let n = Array.length data_x_y_yerr in
    solve n start_par ~f:gsl_multi 
;;

let print_parameters_errors names values errors =
    let n_n = Array.length names and n_v = Array.length values in
    let n = min n_n n_v in
    for k = 0 to n - 1 do
        Printf.printf "%s=%f (err=%f); " names.(k) values.(k) errors.(k) 
    done;
    print_newline ()
;;

let fit_data ?(start_par=[|1.;1.|]) ~n ~f data_x_y_yerr =
    let start_par = if n = Array.length start_par
    then start_par
    else Array.create n 1. in
    let r,_ = fit_lsq ~data_x_y_yerr ~by_function:f ~start_par ~deriv:None in
    let (chi2,ndf) = calc_chi2 data_x_y_yerr (f r) in
    (r,chi2,ndf)
;;

module Error = struct
let calc_cov ~data_x_y_yerr ~f ~deriv ~par covar par_names =
    let chi2,ndf = calc_chi2 data_x_y_yerr (f par) in
    let p = Array.length par in
    let c = max 1. (sqrt (chi2/.(float (ndf-p)))) in
    Array.mapi (fun i _ -> 
        c *. (sqrt covar.(i).(i))
    ) par

open Gsl_root.Polish
let calc_eq ~data_x_y_yerr ~f ~deriv ~par =
(*    print_endline "starting1";*)
  let chi2_fdf n_par = construct_chi2 ~data_x_y_yerr ~f ~deriv ~start:par n_par in
  let make_s n_par = make NEWTON (chi2_fdf n_par) (par.(n_par) *. 1.01) in

  let max_iter = 100 in
(*    print_endline "starting2";*)
    let  a= Array.mapi (fun n_par p ->
      let s = make_s n_par in
      let rec proc i x0 = function
        | true -> x0
        | false when i >= max_iter -> x0
        | false ->
(*    print_endline "starting3";*)
            iterate s ;
            let x = root s in
            let status = Gsl_root.test_delta x x0 0. 1e-3 in
       (*     if status
            then Printf.printf "Converged:\n" ;
            Printf.printf "%5d %10.7f %+10.7f %10.7f\n"
              i x (x -. r_expected) (x -. x0) ;*)
            proc (succ i) x status
      in
      try ((proc 1 p false) -. p)
(*      with _ -> nan*)
      with _ -> 100.
  ) par
    in
(*    print_endline "ending";*)
    a
end

let fit_hist ~by_function:f ~deriv ~start_par hist par_names =
    let data_x_y_yerr' = Array.mapi (fun i y ->
        let x = H1.x_of_bin hist i in
        let error = match hist.H1.sq_errors with
                Some a -> sqrt a.(i)
              | None when y > 0. -> sqrt(y)
              | None -> nan
        in
        (x,y,error)
    ) hist.H1.arr in
    let data_x_y_yerr = Array.of_list (
        List.filter (fun (_,_,e) -> e > 0.) (Array.to_list data_x_y_yerr')) in
    if Array.length data_x_y_yerr = 0
    then (let c () = Array.map (fun _ -> nan) start_par in c(),c())
    else (
        let res, covar = fit_lsq ~data_x_y_yerr ~by_function:f 
                ~deriv ~start_par in
(*        let errors = Error.calc_eq ~data_x_y_yerr ~f ~deriv ~par:res in*)
	let errors = Error.calc_cov ~data_x_y_yerr ~f ~deriv ~par:res covar par_names in

        print_chi2 data_x_y_yerr (f res);
        flush stdout;
        res,errors
    )

let get_fit ~f ?(deriv=None) ?(info="f(x)")
        ~(start: H1.t -> (string * float) array) histogram =
    let hist = get_1d histogram in
    let par_names, start_par = 
        (fun pars -> Array.map fst pars, Array.map snd pars) (start hist) in
    let res,err = fit_hist ~by_function:f ~deriv ~start_par hist par_names in
    print_endline ("fitting by function: "^info);
    print_parameters_errors par_names res err;
   (* info,*) res,err,f,start_par

let get_fit_plot ~f (res,err) histogram =
    let hist = get_1d histogram in
    let h=H1.analog ~n_min:100 ~factor:2 hist in
    init_hist h (f res);
    (H1d h)

let fgaus par x=
    (*par.(0) *. exp(-.(x-.par.(1))**2. /. (par.(2))**2.)*)
    par.(0) *. exp(-.((x-.par.(1))**2.) /. (2. *. (par.(2))**2.))
;;

let fgaus_deriv par xn ~by:l =
  let _A = par.(0) and d = xn -. par.(1) and s = 2. *. (par.(2)**2.) in
  ( match l with 
      | 0 -> exp( -.(d**2.) /. s )
      | 1 -> _A *. exp( -.(d**2.) /. s ) *. (-2.) *. d /. s *. (-1.)
      | 2 -> _A *. exp( -.(d**2.) /. s ) *. 
             ( -.(d**2.)) *. (-1.) /. (s**2.) *. (4. *. par.(2))
  )

let get_fit_gaus =
    get_fit ~f:fgaus ~deriv:(Some fgaus_deriv) 
        ~info:"A*exp{-(x-mean)^2/(2*sigma^2)}"
        ~start:(fun h ->
            let (max,mean,rms_) = H1.max_mean_rms h in
            let rms = if Glob.not_nan rms_ then rms_ 
            else (H1.range h) /. (sqrt 12.)  in
            [| "A", max; "mean", mean; "sigma", rms |])

let fit_peak_by_gaus =
    get_fit ~f:fgaus ~deriv:(Some fgaus_deriv) 
        ~info:"A*exp{-(x-mean)^2/(2*sigma^2)}"
        ~start:(fun h ->
            let (max,mean,rms_) = H1.max_mean_rms h in
            let rms = if Glob.not_nan rms_ then rms_ 
            else (H1.range h) /. (sqrt 12.)  in
            let x = H1.x_of_max h in
            [| "A", max; "mean", x; "sigma", 0.1*.rms |])

let f_linear par x =
    par.(0) +. par.(1) *. x
;;
let f_linear_deriv par x ~by:l =
  match l with 
      0 -> 1. 
    | 1 -> x

let get_fit_linear =
    get_fit ~f:f_linear ~deriv:(Some f_linear_deriv) 
        ~info:"p_0 + p_1 * x"
        ~start:(fun h ->
            let (max,_,_) = H1.max_mean_rms h in
            [| "p_0",0.5*.max; "p_1",0. |] )

let f_exp p x =
    p.(0) *. (exp (p.(1)*.x))

let f_exp_deriv p x ~by:p_l =
    match p_l with
        0 -> exp (p.(1)*.x)
      | 1 -> x *. (f_exp p x)

let get_fit_exp =
    get_fit ~f:f_exp ~deriv:(Some f_exp_deriv)
        ~info:"p0*exp{p1 x}"
        ~start:(fun h ->
            let (_,max) = Glob.minmax h.H1.arr in
            [| "p0",max; "p2",-0. |] )

let f_exp_cy p x =
    p.(0) +. p.(1) *. (exp (p.(2)*.x))

let f_exp_cy_deriv p x ~by:p_l =
    match p_l with
        0 -> 1.
      | 1 -> exp (p.(2)*.x)
      | 2 -> x *. p.(1) *. (exp (p.(2)*.x))

let get_fit_exp_cy =
    get_fit ~f:f_exp_cy ~deriv:(Some f_exp_cy_deriv)
        ~info:"p0 + p1*exp{p2 x}"
        ~start:(fun h ->
            let (min,max) = Glob.minmax h.H1.arr in
            [| "p0",min; "p1",max; "p2",-0. |] )

let f_x2_exp p x =
    p.(0) *. (x**2.) *. (exp (p.(1)*.x))

let f_x2_exp_deriv p x ~by:p_l =
    match p_l with
      0 -> (x**2.) *. (exp (x/.p.(1)))
    | 1 -> x *. (f_x2_exp p x)

let get_fit_x2_exp =
    get_fit ~f:f_x2_exp 
    ~deriv:(Some f_x2_exp_deriv) 
    (*~deriv:None *)
        ~info:"p_0 *x^2 * exp{p_1 * x}"
        ~start:(fun h ->
            let (max,_,_) = H1.max_mean_rms h in
            [| "p_0",0.5*.max; "p_1",-1. |] )

let sgn x = if x < 0. then 0. else 1.

let f_sx2_exp p x =
    let d = x -. p.(2) in
    (sgn d) *. (f_x2_exp p d)

let f_sx2_exp_deriv p x ~by:p_l =
    let d = x -. p.(2) in
    match p_l with
      0 | 1 -> 
          (sgn d) *. (f_x2_exp_deriv p d ~by:p_l)
    | 2 ->
          (sgn d) *. p.(0) *. (exp (p.(1)*.d)) *. (2.*.d +. p.(1)*.(d**2.)) *. (-1.)

let get_fit_sx2_exp =
    get_fit ~f:f_sx2_exp 
        (*~deriv:(Some f_sx2_exp_deriv) *)
        ~deriv:None
        ~info:"p_0 * sgn(x-m) * (x-m)^2 * exp{p_1 * (x-m)}"
        ~start:(fun h ->
            let (max,mean,_) = H1.max_mean_rms h in
            let x = H1.x_of_max h in
            [| "p_0",0.5*.max; "p_1",-1.; "shift", x |] )
           (* [| "p_0",0.5*.max; "p_1",-1.; "shift", mean |] )*)

let f_sqrt_stat p x =
    sqrt (p.(0)**2. +. p.(1)**2./.x)

let get_fit_sqrt_stat =
    get_fit ~f:f_sqrt_stat
        ~deriv:None
        ~info:"sqrt(p0^2 + p1^2/x)"
        ~start:(fun h ->
            let (min,_) = Glob.minmax h.H1.arr in
            [| "p0",min; "p1", 1. |] )

let f_sqrt_sum p x =
    p.(0) +. p.(1) /. sqrt(x)

let get_fit_sqrt_sum =
    get_fit ~f:f_sqrt_sum
        ~deriv:None
        ~info:"p0+p1/sqrt(x)"
        ~start:(fun h ->
            let (min,_) = Glob.minmax h.H1.arr in
            [| "p0",min; "p1", 1. |] )

let f_p0p1p12 p x =
    p.(0) +. p.(1) *. x +. p.(2) /. sqrt(x)

let get_fit_p0p1p12 =
    get_fit ~f:f_p0p1p12
        ~deriv:None
        ~info:"p0+p1 x + p2/sqrt(x)"
        ~start:(fun h ->
            let (min,_) = Glob.minmax h.H1.arr in
            [| "p0",min; "p1", 1.; "p2", 1. |] )

let f_hyp p x =
    p.(0) +. p.(1) /. x

let get_fit_hyp =
    get_fit ~f:f_hyp
        ~deriv:None
        ~info:"p0+p1/x"
        ~start:(fun h ->
            let (min,_) = Glob.minmax h.H1.arr in
            [| "p0",min; "p1", 1. |] )

let f_hyp_s p x =
    p.(0) /. (x-.p.(1))

let get_fit_hyp_s =
    get_fit ~f:f_hyp_s
        ~deriv:None
        ~info:"p0/(x-p1)"
        ~start:(fun h ->
            let (min,max) = Glob.minmax h.H1.arr in
            [| "p0",(max-.min); "p1", 0. |] )

let f_lin_hyp p x =
    p.(0) +. p.(1) *.x +. p.(2) /. x

let get_fit_lin_hyp =
    get_fit ~f:f_lin_hyp
        ~deriv:None
        ~info:"p0+p1 x + p2/x"
        ~start:(fun h ->
            let (min,_) = Glob.minmax h.H1.arr in
            [| "p0",min; "p1", -0.01; "p2", 1000. |] )

(*
let get_fit_gaus_plot histogram =
    let hist = get_1d histogram in
    let (max,mean,rms_) = H1.max_mean_rms hist in
    let rms = if Glob.not_nan rms_ then rms_ 
    else (H1.range hist) /. (sqrt 12.)  in
    let res=fit_gaus ~start_par:[|max;mean;rms|] hist in
    print_parameters [| "A"; "mean"; "sigma" |] res;
    let h=H1.analog ~n_min:100 hist in
    init_hist h (fgaus res);
    (H1d h)
*)

let f_list = [
    "gaussian", get_fit_gaus;
    "gaus.peak", fit_peak_by_gaus;
    "linear", get_fit_linear;
    "exp", get_fit_exp;
    "exp+const", get_fit_exp_cy;
    "^2*exp",get_fit_x2_exp;
    "shifted^2*exp",get_fit_sx2_exp;
    "sqrt stat",get_fit_sqrt_stat;
    "hyp",get_fit_hyp;
    "hyp+lin",get_fit_lin_hyp;
    "p0/(x-p1)",get_fit_hyp_s;
    "p0+p1/sqrt(x)",get_fit_sqrt_sum;
    "p0+p1 x*p2/sqrt(x)",get_fit_p0p1p12;
]

let f_plot_list = 
    List.map (fun (name,fn) -> 
        let get_plot h = 
            let res,err,f,_ = fn h in
            get_fit_plot ~f (res,err) h 
        in
        (name, get_plot)
    ) f_list

